import slangpy as spy
import numpy as np
import pathlib
import torch

device = spy.create_device(
    enable_cuda_interop=True,
    include_paths=[
        pathlib.Path(__file__).parent.absolute(),
    ])


tensor_A = torch.tensor([1.0, 2.0, 3.0], device='cuda', requires_grad=True)
tensor_B = torch.tensor([4.0, 5.0, 6.0], device='cuda', requires_grad=True)

slang_module = spy.TorchModule.load_from_file(
    device=device,
    path='slang/simple_add_kernel.slang')
tensor_C = slang_module.Add(a = tensor_A, b = tensor_B)
loss = torch.sum(tensor_C)
loss.backward()
